Repository navigation
feat(mlops): Add CleanupStep pipeline step - #6468
Conversation
Add a CleanupStep that deletes the Amazon S3 outputs of upstream steps. Each S3CleanupSource names a source step and, for Processing and Transform steps, the exact object URIs to delete. Other rules (supported source types, URI limits, VpcConfig) are left to the service, which validates them at CreatePipeline. depends_on defaults to the source steps since only outputs of completed steps can be deleted. Integ test passed against a production region. --- X-AI-Prompt: Add a PySDK CleanupStep with unit and integration tests and raise a PR X-AI-Tool: kiro-cli
GitHub Actions
Error
GitHub Actions
Error
GitHub Actions
Error
| # The service can only delete outputs of steps that already ran, so the | ||
| # sources are the natural dependencies. | ||
| if depends_on is None: | ||
| depends_on = [source.step for source in s3_sources] |
There was a problem hiding this comment.
shouldn't we add these as additional dependencies if depends_on is not None?
There was a problem hiding this comment.
Good catch, agreed. The service doesn't enforce that a Cleanup step's DependsOn includes its sources, and StepName isn't a property reference, so an explicit depends_on could let Cleanup run before its source finished. The source steps are now always appended to any caller-provided depends_on (deduplicated by name), and the unit test covers the merge.
The service does not require a Cleanup step's DependsOn to include its sources, so an explicit depends_on could let Cleanup run before its source steps finished. --- X-AI-Prompt: Address jkasiraj's PR comment to merge sources into depends_on X-AI-Tool: kiro-cli
GitHub Actions
Error
GitHub Actions
Error
GitHub Actions
Error
| depends_on = [source.step for source in s3_sources] | ||
| # sources are always dependencies, in addition to any the caller passes. | ||
| depends_on = list(depends_on or []) | ||
| dependency_names = {dep if isinstance(dep, str) else dep.name for dep in depends_on} |
There was a problem hiding this comment.
Does this handle all cases? We should add tests that cover each type of case. And use something like:
def _dependency_name(dep: Union[str, Step, StepCollection, StepOutput]) -> str:
"""The name of the step a `depends_on` entry refers to."""
if isinstance(dep, str):
return dep
if isinstance(dep, StepCollection):
return dep.steps[-1].name
if isinstance(dep, StepOutput):
return get_step(dep).name
return dep.name
There was a problem hiding this comment.
Good call, it didn't. A StepOutput entry raised AttributeError (no .name), and a StepCollection was matched by its own name rather than its last step. Switched to _dependency_name as suggested, with parametrized tests over str / Step / StepCollection / StepOutput (043d0c7). Since #6475 merged, I also merged master and taught it StepDependency, so depends_on=[StepDependency(train, TriggerCondition.ANY)] keeps the caller's trigger instead of adding a conflicting Completed edge (b9132bf).
A StepOutput entry raised AttributeError and a StepCollection was matched by its own name instead of its last step, so the source could be appended twice. --- X-AI-Prompt: Address jkasiraj's PR comment to handle all depends_on types X-AI-Tool: kiro-cli
A source listed with a run condition was not recognized, so a second plain Completed edge was appended and compilation failed on conflicting triggers. --- X-AI-Prompt: Fix CleanupStep for the merged run conditions PR X-AI-Tool: kiro-cli
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
🤖 AI code review
This PR adds CleanupStep and S3CleanupSource to sagemaker.mlops.workflow, exposing the Pipelines Cleanup step type, plus the CLEANUP enum value and unit/integ tests. The implementation follows established conventions: source steps are appended (not replacing) to caller-provided depends_on, deduplicated by name while preserving StepDependency triggers; properties raises like FailStep; empty s3_sources/s3_uris are rejected. Fixtures used by the integ test (sagemaker_session, pipeline_session, role, sklearn_latest_version) all resolve from tests/integ/conftest.py. The change is well-tested and I found no substantive issues.
[]
Reviewed commit b9132bf. Automated review; verify before acting.
What: Add CleanupStep and S3CleanupSource to sagemaker.mlops.workflow. A Cleanup step deletes the S3 outputs of upstream steps: StepName-only for Training/Athena sources, StepName + exact S3Uris for Processing/Transform sources. depends_on defaults to the source steps.
Why: SageMaker Pipelines adds a Cleanup step type; this exposes it in the PySDK. Source-type rules, URI limits and VpcConfig checks are left to the service, which validates them at CreatePipeline.
Test:
7 unit tests, including a full Pipeline.definition() shape check; all 366 tests in sagemaker-mlops/tests/unit/workflow pass.
Integ test test_cleanup_step_deletes_processing_output passed in us-west-2: a Processing step writes two objects, Cleanup deletes one, and the test asserts Metadata.Cleanup.CleanedSteps and the resulting S3 state.